import os
import torch
import random
import numpy as np
import pandas as pd
import torch.nn as nn
from tqdm import tqdm
from typing import List, Dict
from torch.optim import Adam
import torch.nn.functional as F
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from dataclasses import dataclass, field
from torch_scatter import scatter_sum, scatter_softmax
from Bio import SeqIO  # pip install biopython


# Define function to read FASTA files using Biopython
def read_fasta_biopython(file_path):
    sequences = {}
    for record in SeqIO.parse(file_path, "fasta"):
        sequences[record.id] = str(record.seq)
    return sequences


# Define RNADataset Class and Seeding Function
class RNADatasetV2(Dataset):
    def __init__(self, data_path, is_train=True):
        super(RNADatasetV2, self).__init__()
        self.npy_dir = data_path + "/coords"
        self.name_list = [i[:-4] for i in os.listdir(data_path + "/coords")]
        self.seq_dir =  data_path + "/seqs/"
        self.cache = {}
        self.is_train = is_train

    def __len__(self):
        return len(self.name_list)

    def get_pdb_id(self, idx):
        return self.name_list[idx]

    def merge_coords_file_path(self, pdb_id):
        return os.path.join(self.npy_dir, pdb_id + '.npy')

    def load_feature(self, pdb_id):
        coords = np.load(self.merge_coords_file_path(pdb_id))
        feature = {
            "name": pdb_id,
            "coords": {
                "P": coords[:, 0, :],
                "O5'": coords[:, 1, :],
                "C5'": coords[:, 2, :],
                "C4'": coords[:, 3, :],
                "C3'": coords[:, 4, :],
                "O3'": coords[:, 5, :],
                "N": coords[:, 6, :],
            }
        }

        return feature
    
    def load_seq(self, pdb_id):
        return list(read_fasta_biopython(self.seq_dir + pdb_id + ".fasta").values())[0]
    
    def first_load(self, idx):
        pdb_id = self.get_pdb_id(idx)
        feature = self.load_feature(pdb_id)
        if self.is_train:
            feature["seq"] = self.load_seq(pdb_id)
        else:
            feature["seq"] = None
        return feature

    def __getitem__(self, idx):
        if idx in self.cache:
            return self.cache[idx]
        else:
            data = self.first_load(idx)
            self.cache[idx] = data
            return data
    
    def __iter__(self):
        for idx in range(len(self)):
            yield self.__getitem__(idx)

    def get_lengths_by_indices(self, indices):
        lengths = []
        for idx in indices:
            pdb_id = self.get_pdb_id(idx)
            file_path = self.merge_coords_file_path(pdb_id)
            with open(file_path, 'rb') as f:
                # 读取文件头的前8字节（魔数和版本号）
                version = np.lib.format.read_magic(f)
                # 读取文件头信息（包含shape/dtype等）
                shape, _, _ = np.lib.format._read_array_header(f, version)
            lengths.append(shape[0])
        return lengths
    
def prepare_rna_batch(rna_samples):
    """将RNA样本批次转换为神经网络可用的张量格式
    
    参数:
        rna_samples: 包含多个RNA样本的列表，每个样本应包含:
            - 'seq': 碱基序列字符串(如"AUCG")
            - 'coords': 原子坐标字典
            - 'name': 样本标识名
    
    返回:
        features: 形状[B, L_max, 6, 3]的原子坐标张量
        labels: 形状[B, L_max]的碱基类别张量  
        valid_mask: 有效位置掩码
        seq_lengths: 各样本实际长度
        sample_names: 样本名称列表
    """
    tag_list = list("AUCG")
    idx2tag = {i: tag for i, tag in enumerate(tag_list)}
    tag2idx = {tag: i for i, tag in enumerate(tag_list)}
    # KEY_ATOMS = ["P", "O5'", "C5'", "C4'", "C3'", "O3'", "N"]
    KEY_ATOMS = ["P", "O5'", "C5'", "C4'", "C3'", "O3'"]
    
    # 初始化维度
    batch_size = len(rna_samples)
    seq_lengths = [b["coords"]["P"].shape[0] for b in rna_samples]  # 获取RNA序列长度
    max_len = max(seq_lengths)  # 当前batch最大长度
    
    # 预分配张量（直接初始化为最终类型）
    coords = np.zeros((batch_size, max_len, 6, 3), dtype=np.float32)  # batch, 序列长度，6个原子位置，每个原子坐标
    labels = np.full((batch_size, max_len), -1, dtype=np.int64)
    mask = np.zeros((batch_size, max_len), dtype=np.float32)
    
    for i, sample in enumerate(rna_samples):
        curr_len = seq_lengths[i]
        # 向量化处理原子坐标
        atom_stack = np.stack([np.nan_to_num(sample['coords'][atom], nan=0.0) for atom in KEY_ATOMS],
                              axis=1)  # N 6 3
        # 有效位置标记
        valid_pos = np.isfinite(atom_stack.sum((1,2)))
        actual_len = min(curr_len, sum(valid_pos))
        # 单样本赋值
        coords[i,:actual_len] = atom_stack[valid_pos][:actual_len]
        mask[i,:actual_len] = 1.0
        
        # 有标签的话读取标签
        seq = np.array([c for c in sample['seq']], dtype=np.str_)  # 字符串转列表再转numpy
        labels[i,:actual_len] = [tag2idx[b] for b in seq[valid_pos][:actual_len]]
    
    features, labels, valid_mask, seq_lengths, sample_names = (torch.from_numpy(coords),
        torch.from_numpy(labels),
        torch.from_numpy(mask),
        np.array(seq_lengths),
        [s['name'] for s in rna_samples])
    
    return features, labels, valid_mask, seq_lengths, sample_names
